#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
template<typename TYPE_X,typename TYPE_Y> class KernelReplicationPad2d {
    using T=TYPE_X;
public:
    __aicore__ inline KernelReplicationPad2d() {}
    // __aicore__ inline void Init(GM_ADDR x, GM_ADDR paddings,  GM_ADDR y, int32_t M, int32_t N, int32_t Z, int32_t dim1,int32_t dim2,
    // int32_t dim3,int32_t dim4,int32_t new_y_nums) {
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR paddings,  GM_ADDR y, int32_t M, int32_t N, int32_t Z,int32_t pad_dim1,int32_t pad_dim2,int32_t pad_dim3,int32_t pad_shapeNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        paddingsGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_PADDINGS *>(paddings), pad_dim1);
        // paddingsGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_PADDINGS *>(paddings), pad_dim2);
        // paddingsGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_PADDINGS *>(paddings), pad_dim3);
        this->M = M; //行
        this->N = N; //列
        this->Z = Z; //循环次数
        // this->dim1 =dim1;//Pading的第一个维度  左
        // this->dim2 =dim2;//Pading的第二个维度  右
        // this->dim3 =dim3;//Pading的第三个维度  上
        // this->dim4 =dim4;//Pading的第四个维度  下
        this->pad_dim1 = pad_dim1;
        this->pad_dim2 = pad_dim2;
        this->pad_dim3 = pad_dim3;
        this->pad_shapeNum = pad_shapeNum;
        this->new_y_nums = new_y_nums;//新生成的y的总维度
        this->row32 = (N+31) /32 *32;

        this->newrow32 =(this->dim1+this->dim2+N+31)/32*32;
        // this->dim_size = 1;
        // for(int32_t i = 0;i < pad_dim; ++i){
        //     this->dim[i] = paddingsGm.GetValue(i);
        // }
        // this->dim = this->dim[i];
        if(pad_shapeNum == 1){

            if(this->pad_dim1==4){
                this->dim1 = paddingsGm.GetValue(this->pad_dim1-4);
                this->dim2 = paddingsGm.GetValue(this->pad_dim1-3);
                this->dim3 = paddingsGm.GetValue(this->pad_dim1-2);
                this->dim4 = paddingsGm.GetValue(this->pad_dim1-1);

            }
            else if(this->pad_dim1==1){

                this->dim1 = paddingsGm.GetValue(this->pad_dim1-1);
                this->dim2 = paddingsGm.GetValue(this->pad_dim1-1);
                this->dim3 = paddingsGm.GetValue(this->pad_dim1-1);
                this->dim4 = paddingsGm.GetValue(this->pad_dim1-1);

            }
            else if(this->pad_dim1==2){
                this->dim1 = paddingsGm.GetValue(this->pad_dim1-2);
                this->dim2 = paddingsGm.GetValue(this->pad_dim1-1);
                this->dim3 = 0;
                this->dim4 = 0;


            }
            
        }
        else if(pad_shapeNum == 2){
            this->dim1 = paddingsGm.GetValue(this->pad_dim1-2);
            this->dim2 = paddingsGm.GetValue(this->pad_dim1-1);
            this->dim3 = paddingsGm.GetValue(this->pad_dim2);
            this->dim4 = paddingsGm.GetValue(this->pad_dim2+1);
        }else if(pad_shapeNum == 3){
            this->dim1 = paddingsGm.GetValue(this->pad_dim1*this->pad_dim2*this->pad_dim3-4);
            this->dim2 = paddingsGm.GetValue(this->pad_dim1*this->pad_dim2*this->pad_dim3-3);
            this->dim3 = paddingsGm.GetValue(this->pad_dim1*this->pad_dim2*this->pad_dim3-2);
            this->dim4 = paddingsGm.GetValue(this->pad_dim1*this->pad_dim2*this->pad_dim3-1);
        }
        // this->dim1 = paddingsGm.GetValue(this->pad_dim-4);
        // this->dim2 = paddingsGm.GetValue(this->pad_dim-3);
        // this->dim3 = paddingsGm.GetValue(this->pad_dim-2);
        // this->dim4 = paddingsGm.GetValue(this->pad_dim-1);
        this->new_y_nums = Z *(M+this->dim1+this->dim2)*(N+this->dim3+this->dim4); 
        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(x), M * N * Z);
        yGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_Y *>(y), new_y_nums);
        //zGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_Y *>(y), (N+this->dim1+this->dim2));
        //pipe.InitBuffer(inQueueX, BUFFER_NUM, (this->dim1+N+31+this->dim2)/32*32* sizeof(DTYPE_X));

        //pipe.InitBuffer(inQueueZ, 1, (this->dim1+N+31+this->dim2)/32*32* sizeof(DTYPE_X));


        //pipe.InitBuffer(outQueueY, BUFFER_NUM, (this->dim1+this->dim2+N+31)/32*32 * sizeof(DTYPE_Y));

        //pipe.InitBuffer(QueueTemp, dim1 * sizeof(DTYPE_X));

    }
    __aicore__ inline void Process() {



        for(int i=0;i<new_y_nums;i++)
        {
            yGm.SetValue(i,0.0);
        }
         for(int32_t z1 = 0; z1 < Z; z1++){
            for (int32_t a=0; a<dim3;a++){   // a = 0 ,1   dim3=2
                
                for(int32_t i=0; i<dim1;i++){
                     yGm.SetValue(z1*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2)+a*(this->dim1+N+this->dim2)+i,xGm.GetValue(z1*(M)*(N)));
            }   
                for(int32_t j=dim1;j<N+dim1;j++){

                    yGm.SetValue(z1*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2)+a*(this->dim1+N+this->dim2)+j,xGm.GetValue(z1*(M)*(N)+j-dim1));
                }

                for(int32_t z=N+dim1;z<N+dim1+dim2;z++){

                    yGm.SetValue(z1*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2)+a*(this->dim1+N+this->dim2)+z,xGm.GetValue(z1*(M)*(N)+N-1));
                } 
            
        }   
            for(int32_t b=dim3;b<dim3+N;b++){  // b=2

                 for(int32_t i=0; i<dim1;i++){
                     yGm.SetValue(z1*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2)+b*(this->dim1+N+this->dim2)+i,xGm.GetValue(z1*(M)*(N)+(b-dim3)*N));
            }   
                for(int32_t j=dim1;j<N+dim1;j++){

                    yGm.SetValue(z1*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2)+b*(this->dim1+N+this->dim2)+j,xGm.GetValue(z1*(M)*(N)+(b-dim3)*N+j-dim1));
                }

                for(int32_t z=N+dim1;z<N+dim1+dim2;z++){

                    yGm.SetValue(z1*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2)+b*(this->dim1+N+this->dim2)+z,xGm.GetValue(z1*(M)*(N)+(b-dim3)*N+N-1));
                } 



            }

            for(int32_t c=dim3+M;c<dim3+dim4+M;c++){

                for(int32_t i=0; i<dim1;i++){
                     yGm.SetValue(z1*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2)+c*(this->dim1+N+this->dim2)+i,xGm.GetValue(z1*(M)*(N)+(M-1)*N));
            }   
                for(int32_t j=dim1;j<N+dim1;j++){

                    yGm.SetValue(z1*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2)+c*(this->dim1+N+this->dim2)+j,xGm.GetValue(z1*(M)*(N)+(M-1)*N+j-dim1));
                }

                for(int32_t z=N+dim1;z<N+dim1+dim2;z++){

                    yGm.SetValue(z1*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2)+c*(this->dim1+N+this->dim2)+z,xGm.GetValue(z1*(M)*(N)+(M-1)*N+N-1));
                } 


            }
        
        }



        //  for(int32_t i = 0; i < Z; i ++)
        // {  
            // for(int32_t j = 0; j < M; j++)  //所有行
            // {
                // CopyIn(i,j);
                // Compute(i,j);
                // CopyOut(i,j);
                
            //}
        //}


        // for(int32_t i = 0; i < Z; i ++)
        // {  
        //     for(int32_t j = 0; j < M; j++)
        //     {
        //         CopyIn(i,j);
        //         Compute(i,j);
        //         CopyOut(i,j);
        //     }
        // }
    }
private:
     __aicore__ inline void CopyIn(int32_t i, int32_t j)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();

        if(this->dim1 <=8 && this->dim2 <=8){   

            
            DataCopyExtParams copyParams{1, static_cast<uint32_t>(N * sizeof(DTYPE_X)), 0, 0, 0};

            DataCopyPadExtParams<DTYPE_X> padParams{true, static_cast<uint8_t>(this->dim1), static_cast<uint8_t>(this->dim2),  static_cast<DTYPE_X>(xGm.GetValue(i*M*N+j*N+(N-1)))};
            // DataCopyPadExtParams<DTYPE_Y> padParams2{true, static_cast<uint8_t>(dim1), 0,  static_cast<DTYPE_X>(x1)};
            // // DataCopyPadExtParams<DTYPE_Y> padParams{true, dim1, dim2,  static_cast<DTYPE_X>(xn)};
            DataCopyPad(xLocal,xGm[i*M*N+j*N],copyParams,padParams);

            DTYPE_X x1= xLocal.GetValue(this->dim1); //第一行的第一个元素
        
            //DTYPE_X xn =xLocal.GetValue((dim1+N-1));
            // for(int32_t i=0;i<dim1;i++){
                
            //     xLocal.SetValue(i,x1);

            // }
            Duplicate(xLocal,x1,this->dim1);
            //DataCopy(xLocal, xGm[i*M*N + j*N], row32);
            inQueueX.EnQue(xLocal);

        }

         else  {

            LocalTensor<DTYPE_X> zLocal = inQueueZ.AllocTensor<DTYPE_X>();



            for(int i=0;i<(new_y_nums);i++){

                yGm.SetValue(i,0.0);
            }  

            if  constexpr (std::is_same_v<DTYPE_X,float>){

                if(this->dim1 >8 && this->dim2 <=8){ //假设this->dim1=9

                    int  loopcount= this->dim1 % 8;

                    DataCopyExtParams copyParams{1, static_cast<uint32_t>(N * sizeof(DTYPE_X)), 0, 0, 0};

                    DataCopyPadExtParams<DTYPE_X> padParams{true, static_cast<uint8_t>(8), static_cast<uint8_t>(this->dim2),  static_cast<DTYPE_X>(xGm.GetValue(i*M*N+j*N+(N-1)))};
                        // DataCopyPadExtParams<DTYPE_Y> padParams2{true, static_cast<uint8_t>(dim1), 0,  static_cast<DTYPE_X>(x1)};
                        // // DataCopyPadExtParams<DTYPE_Y> padParams{true, dim1, dim2,  static_cast<DTYPE_X>(xn)};
                    DataCopyPad(xLocal,xGm[i*M*N+j*N],copyParams,padParams);


                    DTYPE_X x1= xLocal.GetValue(8); //第一行的第一个元素
        
                    Duplicate(xLocal,x1,8);



                    DataCopy(yGm,xLocal,this->newrow32);

                    for(int i=0;i<loopcount-1;i++){  //假设loopcount=2 this->dim1=17

                        DataCopyExtParams copyParams{1, static_cast<uint32_t>((8*(i+1)+N+this->dim2) * sizeof(DTYPE_X)), 0, 0, 0};

                        DataCopyPadExtParams<DTYPE_X> padParams{true, static_cast<uint8_t>(8), static_cast<uint8_t>(0),  static_cast<DTYPE_X>(yGm.GetValue(0))};
                            // DataCopyPadExtParams<DTYPE_Y> padParams2{true, static_cast<uint8_t>(dim1), 0,  static_cast<DTYPE_X>(x1)};
                            // // DataCopyPadExtParams<DTYPE_Y> padParams{true, dim1, dim2,  static_cast<DTYPE_X>(xn)};
                        DataCopyPad(xLocal,yGm,copyParams,padParams);


                        for(int i=0;i<(new_y_nums);i++){
                            yGm.SetValue(i,0.0);
                     }  

                        DataCopy(yGm,xLocal,this->newrow32);

                    }

                    uint8_t tailnum = this->dim1-(loopcount*8);

                    if(tailnum !=0){

                        DataCopyExtParams copyParams{1, static_cast<uint32_t>((8*loopcount+N+this->dim2) * sizeof(DTYPE_X)), 0, 0, 0};

                        DataCopyPadExtParams<DTYPE_X> padParams{true, static_cast<uint8_t>(tailnum), static_cast<uint8_t>(0),  static_cast<DTYPE_X>(yGm.GetValue(0))};
                            // DataCopyPadExtParams<DTYPE_Y> padParams2{true, static_cast<uint8_t>(dim1), 0,  static_cast<DTYPE_X>(x1)};
                            // // DataCopyPadExtParams<DTYPE_Y> padParams{true, dim1, dim2,  static_cast<DTYPE_X>(xn)};
                        DataCopyPad(zLocal,yGm,copyParams,padParams);

                        Copy(xLocal,zLocal,(this->dim1+this->dim2+N),1,{1,1,8,8});

                        inQueueZ.EnQue(zLocal);

                        LocalTensor<DTYPE_X> z1Local = inQueueZ.DeQue<DTYPE_X>();

                        inQueueZ.FreeTensor(z1Local);







                        // DataCopy(zGm,xLocal,this->newrow32);
                        

                    }

                    for(int i=0;i<(new_y_nums);i++){
                            yGm.SetValue(i,0.0);
                     }  

                    inQueueX.EnQue(xLocal);

                    



                        // DTYPE_X x1= xLocal.GetValue(8); //第一行的第一个元素
            
                        // Duplicate(xLocal,x1,this->dim1);



                    //DataCopy(zGm,xLocal,this->newrow32);


                    
        
                    //inQueueX.EnQue(xLocal);

               

                        


                    //}



                }

                if(this->dim1 <=8 && this->dim2 >8){

                    int a=0;


                }
                
                if(this->dim1 >8 && this->dim2 >8){

                    int b=0;


            }


            }

            else{

                 if(this->dim1 >16 && this->dim2 <=16){
                    int c=0;


                }

                if(this->dim1 <=16 && this->dim2 >=16){
                    int d=0;


                }
                
                if(this->dim1 >16&& this->dim2 >16){

                    int e=0;


            }



            }

            


         }
        //
        // Duplicate(xLocal,x1,dim1);
        // 
        //inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void Compute(int32_t i, int32_t j)
    {

        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();

        //LocalTensor<DTYPE_X> p1 = QueueTemp.Get<DTYPE_X>();

        Adds(yLocal,xLocal,static_cast<DTYPE_X>(0),(this->dim1+N+this->dim2));

        outQueueY.EnQue<DTYPE_Y>(yLocal);

        inQueueX.FreeTensor(xLocal);

        // 

        //DataCopyExtParams copyParams{1, static_cast<uint32_t>(N * sizeof(DTYPE_X)), 0, 0, 0};

        //DataCopyPadExtParams<DTYPE_Y> padParams{true, static_cast<uint8_t>(dim1), static_cast<uint8_t>(dim2),  static_cast<DTYPE_X>(xn)};
        // DataCopyPadExtParams<DTYPE_Y> padParams2{true, static_cast<uint8_t>(dim1), 0,  static_cast<DTYPE_X>(x1)};
        // // DataCopyPadExtParams<DTYPE_Y> padParams{true, dim1, dim2,  static_cast<DTYPE_X>(xn)};
        //DataCopyPad(yLocal,xGm[i*M*N+j*N],copyParams,padParams);
        // DataCopyPad(yLocal2,yLocal1,copyParams,padParams);
        //Duplicate(yLocal,x1,dim1);

        // Sub(yLocal,yLocal,yLocal,dim1);

        // Adds(yLocal,yLocal,x1,dim1);

        // for(int32_t i=0;i<dim1;i++){

        //     yLocal.SetValue(i,x1);
        // }

        // Duplicate(p1,x1,dim1);

        // Copy(yLocal,p1,dim1,1,{1,1,8,8});

        //LocalTensor<DTYPE_Y> yLocal1 = outQueueY.DeQue<DTYPE_Y>();


        // DataCopyExtParams copyParams{1, static_cast<uint32_t>(N * sizeof(DTYPE_X)), 0, 0, 0};

        // DataCopyPadExtParams<DTYPE_Y> padParams{true, static_cast<uint8_t>(dim1), 0,  static_cast<DTYPE_Y>(x1)};
        // DataCopyPad(yLocal,xGm[i*M*N+j*N],copyParams,padParams);
    }
    __aicore__ inline void CopyOut(int32_t i, int32_t j)
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();

         if(j>=1 && j<(M-1)){
            DataCopyExtParams copyParams{1, static_cast<uint32_t>((this->dim1+this->dim2+N) * sizeof(DTYPE_X)), 0, 0, 0};
            DataCopyPad(yGm[i*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2) + (j+this->dim3)*(this->dim1+N+this->dim2)], yLocal, copyParams); 

            //DataCopy(yGm[i*(dim3+M+dim4)*(dim1+N+dim2) + (j+dim3)*(dim1+N+dim2)], xLocal, row32);

            //yGm.SetValue()
        }

        else if(j==0){

            for(int32_t a=0;a<(this->dim3+1);a++){

                 //DataCopy(yGm[i*(dim3+M+dim4)*(dim1+N+dim2) + (a+j)*(dim1+N+dim2)], xLocal, row32);
                 DataCopyExtParams copyParams{1, static_cast<uint32_t>((this->dim1+this->dim2+N) * sizeof(DTYPE_X)), 0, 0, 0};
                 DataCopyPad(yGm[i*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2) + (j+a)*(this->dim1+N+this->dim2)], yLocal, copyParams); 

            }

            // for(int32_t b=0; b<dim1;b++){

            //  yGm.SetValue((i*(dim3+M+dim4)*(dim1+N+dim2) + b+ (a+j)*(dim1+N+dim2)),x1);

            // }
        }
        else if(j==(M-1)){

            for(int32_t a=0;a<(this->dim4+1);a++){

                 //DataCopy(yGm[i*(dim3+M+dim4)*(dim1+N+dim2) + (a+j+dim3)*(dim1+N+dim2)], xLocal, row32);

                 DataCopyExtParams copyParams{1, static_cast<uint32_t>((this->dim1+this->dim2+N) * sizeof(DTYPE_X)), 0, 0, 0};
                 DataCopyPad(yGm[i*(this->dim3+M+this->dim4)*(this->dim1+N+this->dim2) + (a+j+this->dim3)*(this->dim1+N+this->dim2)], yLocal, copyParams); 
 
            }

        }
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECIN, 1> inQueueZ;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    //TBuf<QuePosition::VECCALC> QueueTemp;
    GlobalTensor<DTYPE_X> xGm;
    //GlobalTensor<DTYPE_PADDINGS> pGm;
    GlobalTensor<DTYPE_Y> yGm;

    GlobalTensor<DTYPE_Y> zGm;

    GlobalTensor<DTYPE_PADDINGS> paddingsGm;
    DataCopyPadExtParams<DTYPE_X> padParams;
    DataCopyExtParams copyParams;
    int32_t M;
    int32_t N;
    int32_t Z;
    int32_t dim1;
    int32_t dim2;
    int32_t dim3;
    int32_t dim4;
    int32_t pad_dim1;
    int32_t pad_dim2;
    int32_t pad_dim3;
    int32_t pad_shapeNum;
    int32_t new_y_nums;
    int32_t row32;
    int32_t newrow32;
};


extern "C" __global__ __aicore__ void replication_pad2d(GM_ADDR x, GM_ADDR paddings, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelReplicationPad2d<DTYPE_X,DTYPE_Y> op;
    // op.Init(x, paddings, y,tiling_data.M, tiling_data.N, tiling_data.Z,
    // tiling_data.dim1, tiling_data.dim2, tiling_data.dim3,tiling_data.dim4, tiling_data.new_y_nums);
    op.Init(x, paddings, y,tiling_data.M, tiling_data.N, tiling_data.Z,tiling_data.pad_dim1,tiling_data.pad_dim2,tiling_data.pad_dim3,tiling_data.pad_shapeNum);
    op.Process();
}
